#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
图像Mask生成脚本

功能：
1. 遍历指定目录中的所有图像文件
2. 查找对应的JSON标注文件
3. 如果JSON不存在，生成纯黑的二值图
4. 如果JSON存在，根据标签顺序递增生成对应的mask图像
5. 同时保存原图像和mask图像到不同的子目录中

使用方法：
# 使用标签顺序（递增映射）
python utils/mask_generator.py --input_dir /path/to/images --output_dir /path/to/masks --label_order "dark,light,edge,defect"

# 使用标签合并配置（自定义映射）
python utils/mask_generator.py --input_dir /path/to/images --output_dir /path/to/masks --label_merge_config "dark:1,mohu:1,light:2,edge:3"

注意：背景默认为0，不需要传入。使用--label_order时，传入的标签从1开始递增；使用--label_merge_config时，可以自定义标签到值的映射，支持多个标签映射到同一个值。
"""

import os
import json
import argparse
from pathlib import Path
from typing import List, Dict, Any, Optional, Tuple
import numpy as np
from PIL import Image, ImageDraw
import cv2


class MaskGenerator:
    """图像Mask生成器"""
    
    def __init__(self, label_order: List[str] = None, label_merge_config: Dict[str, int] = None):
        """
        初始化Mask生成器
        
        Args:
            label_order (List[str]): 标签顺序列表，用于生成递增的mask值
                                 背景默认为0，传入的标签从1开始递增
            label_merge_config (Dict[str, int]): 标签合并配置，指定标签到值的映射
                                               如果提供，将覆盖label_order的递增逻辑
        """
        self.label_order = label_order or []
        self.label_merge_config = label_merge_config or {}
        
        # 构建标签到值的映射
        if self.label_merge_config:
            # 使用合并配置
            self.label_to_value = self.label_merge_config.copy()
            print(f"标签映射（合并配置）: 背景=0, {self.label_to_value}")
        else:
            # 使用原来的递增逻辑
            self.label_to_value = {label: idx + 1 for idx, label in enumerate(self.label_order)}
            print(f"标签映射（递增配置）: 背景=0, {self.label_to_value}")
    
    def load_json_annotations(self, json_path: str) -> Optional[Dict[str, Any]]:
        """
        加载JSON标注文件
        
        Args:
            json_path (str): JSON文件路径
            
        Returns:
            Optional[Dict[str, Any]]: 标注数据，如果文件不存在或解析失败则返回None
        """
        if not os.path.exists(json_path):
            return None
        
        try:
            with open(json_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
            return data
        except Exception as e:
            print(f"警告：无法解析JSON文件 {json_path}: {str(e)}")
            return None
    
    def create_black_mask(self, image_size: Tuple[int, int]) -> np.ndarray:
        """
        创建纯黑的二值图
        
        Args:
            image_size (Tuple[int, int]): 图像尺寸 (width, height)
            
        Returns:
            np.ndarray: 纯黑的二值图
        """
        return np.zeros((image_size[1], image_size[0]), dtype=np.uint8)
    
    def create_mask_from_annotations(self, 
                                   annotations: Dict[str, Any], 
                                   image_size: Tuple[int, int]) -> np.ndarray:
        """
        根据标注数据创建mask图像
        
        Args:
            annotations (Dict[str, Any]): 标注数据
            image_size (Tuple[int, int]): 图像尺寸 (width, height)
            
        Returns:
            np.ndarray: mask图像
        """
        # 创建黑色背景
        mask = np.zeros((image_size[1], image_size[0]), dtype=np.uint8)
        
        # 处理不同类型的标注格式
        if 'shapes' in annotations:
            # LabelMe格式
            self._process_labelme_shapes(annotations['shapes'], mask)
        elif 'annotations' in annotations:
            # COCO格式
            self._process_coco_annotations(annotations['annotations'], mask)
        elif 'objects' in annotations:
            # 自定义格式
            self._process_custom_objects(annotations['objects'], mask)
        else:
            # 尝试直接处理标注数据
            self._process_direct_annotations(annotations, mask)
        
        return mask
    
    def _process_labelme_shapes(self, shapes: List[Dict], mask: np.ndarray):
        """处理LabelMe格式的shapes"""
        for shape in shapes:
            label = shape.get('label', '')
            if label not in self.label_to_value:
                continue
            
            value = self.label_to_value[label]
            points = shape.get('points', [])
            
            if shape.get('shape_type') == 'polygon' and len(points) >= 3:
                # 多边形
                pts = np.array(points, dtype=np.int32)
                cv2.fillPoly(mask, [pts], value)
            elif shape.get('shape_type') == 'rectangle' and len(points) == 2:
                # 矩形
                x1, y1 = map(int, points[0])
                x2, y2 = map(int, points[1])
                cv2.rectangle(mask, (x1, y1), (x2, y2), value, -1)
            elif shape.get('shape_type') == 'circle' and len(points) == 2:
                # 圆形
                center = tuple(map(int, points[0]))
                radius = int(np.linalg.norm(np.array(points[1]) - np.array(points[0])))
                cv2.circle(mask, center, radius, value, -1)
    
    def _process_coco_annotations(self, annotations: List[Dict], mask: np.ndarray):
        """处理COCO格式的annotations"""
        for ann in annotations:
            category_id = ann.get('category_id', 0)
            segmentation = ann.get('segmentation', [])
            
            if segmentation:
                # 处理分割多边形
                for seg in segmentation:
                    if len(seg) >= 6:  # 至少需要3个点
                        pts = np.array(seg).reshape(-1, 2).astype(np.int32)
                        cv2.fillPoly(mask, [pts], category_id)
    
    def _process_custom_objects(self, objects: List[Dict], mask: np.ndarray):
        """处理自定义格式的objects"""
        for obj in objects:
            label = obj.get('label', '')
            if label not in self.label_to_value:
                continue
            
            value = self.label_to_value[label]
            bbox = obj.get('bbox', [])
            polygon = obj.get('polygon', [])
            
            if polygon and len(polygon) >= 3:
                # 多边形
                pts = np.array(polygon, dtype=np.int32)
                cv2.fillPoly(mask, [pts], value)
            elif bbox and len(bbox) == 4:
                # 边界框 [x, y, width, height]
                x, y, w, h = map(int, bbox)
                cv2.rectangle(mask, (x, y), (x + w, y + h), value, -1)
    
    def _process_direct_annotations(self, annotations: Dict, mask: np.ndarray):
        """直接处理标注数据"""
        # 尝试处理各种可能的标注格式
        for key, value in annotations.items():
            if isinstance(value, list) and len(value) > 0:
                if key in self.label_to_value:
                    mask_value = self.label_to_value[key]
                    # 假设value是坐标点列表
                    if isinstance(value[0], (list, tuple)) and len(value[0]) == 2:
                        pts = np.array(value, dtype=np.int32)
                        cv2.fillPoly(mask, [pts], mask_value)
    
    def process_single_image(self, 
                           image_path: str, 
                           output_dir: str) -> bool:
        """
        处理单张图像，生成对应的mask并保存原图像
        
        Args:
            image_path (str): 输入图像路径
            output_dir (str): 输出目录路径
            
        Returns:
            bool: 处理是否成功
        """
        try:
            # 加载图像获取尺寸
            image = Image.open(image_path)
            image_size = image.size  # (width, height)
            
            # 获取文件名（不含扩展名）
            image_name = Path(image_path).stem
            image_ext = Path(image_path).suffix
            
            # 查找对应的JSON文件
            json_path = Path(image_path).with_suffix('.json')
            
            if not os.path.exists(json_path):
                # JSON不存在，生成纯黑的二值图
                print(f"  JSON文件不存在: {json_path}")
                mask = self.create_black_mask(image_size)
            else:
                # JSON存在，根据标注生成mask
                print(f"  处理JSON文件: {json_path}")
                annotations = self.load_json_annotations(str(json_path))
                
                if annotations is None:
                    # JSON解析失败，生成纯黑的二值图
                    mask = self.create_black_mask(image_size)
                else:
                    # 根据标注生成mask
                    mask = self.create_mask_from_annotations(annotations, image_size)
            
            # 创建输出目录结构
            img_dir = os.path.join(output_dir, 'img')
            label_dir = os.path.join(output_dir, 'label')
            os.makedirs(img_dir, exist_ok=True)
            os.makedirs(label_dir, exist_ok=True)
            
            # 保存原图像到img目录
            img_output_path = os.path.join(img_dir, f"{image_name}{image_ext}")
            image.save(img_output_path)
            
            # 保存mask图像到label目录
            label_output_path = os.path.join(label_dir, f"{image_name}.png")
            mask_image = Image.fromarray(mask, mode='L')
            mask_image.save(label_output_path)
            
            print(f"  ✓ 保存原图像: {img_output_path}")
            print(f"  ✓ 保存mask: {label_output_path}")
            return True
            
        except Exception as e:
            print(f"  ✗ 处理失败: {str(e)}")
            return False
    
    def process_directory(self, 
                        input_dir: str, 
                        output_dir: str,
                        supported_formats: Tuple[str, ...] = ('.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif')) -> Dict[str, int]:
        """
        批量处理目录中的所有图像
        
        Args:
            input_dir (str): 输入图像目录
            output_dir (str): 输出mask目录
            supported_formats (Tuple[str, ...]): 支持的图像格式
            
        Returns:
            Dict[str, int]: 处理统计信息
        """
        if not os.path.exists(input_dir):
            print(f"✗ 错误：输入目录不存在: {input_dir}")
            return {'total': 0, 'success': 0, 'failed': 0}
        
        # 创建输出目录
        os.makedirs(output_dir, exist_ok=True)
        print(f"✓ 输出目录: {output_dir}")
        
        # 查找所有图像文件
        image_files = []
        for file in os.listdir(input_dir):
            if Path(file).suffix.lower() in supported_formats:
                image_files.append(file)
        
        if not image_files:
            print("✗ 警告：输入目录中没有找到图像文件")
            return {'total': 0, 'success': 0, 'failed': 0}
        
        print(f"✓ 找到 {len(image_files)} 个图像文件")
        
        # 处理统计
        stats = {'total': len(image_files), 'success': 0, 'failed': 0}
        
        # 处理每张图像
        for i, filename in enumerate(image_files, 1):
            print(f"\n处理图像 {i}/{len(image_files)}: {filename}")
            
            input_path = os.path.join(input_dir, filename)
            
            if self.process_single_image(input_path, output_dir):
                stats['success'] += 1
            else:
                stats['failed'] += 1
        
        return stats


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='图像Mask生成工具')
    parser.add_argument('--input_dir', '-i', required=True, help='输入图像目录路径')
    parser.add_argument('--output_dir', '-o', required=True, help='输出mask目录路径')
    parser.add_argument('--label_order', '-l', 
                       help='标签顺序，用逗号分隔，例如: "class1,class2,class3,class4"')
    parser.add_argument('--label_merge_config', '-m',
                       help='标签合并配置，格式: "label1:value1,label2:value2"，例如: "dark:1,mohu:1,light:2"')
    
    args = parser.parse_args()
    
    # 解析标签顺序
    label_order = None
    if args.label_order:
        label_order = [label.strip() for label in args.label_order.split(',')]
    
    # 解析标签合并配置
    label_merge_config = None
    if args.label_merge_config:
        label_merge_config = {}
        for pair in args.label_merge_config.split(','):
            if ':' in pair:
                label, value = pair.strip().split(':', 1)
                label_merge_config[label.strip()] = int(value.strip())
            else:
                print(f"警告：忽略无效的合并配置项: {pair}")
    
    # 检查参数
    if not label_order and not label_merge_config:
        print("错误：必须提供 --label_order 或 --label_merge_config 参数")
        return
    
    print("="*60)
    print("图像Mask生成工具")
    print("="*60)
    print(f"输入目录: {args.input_dir}")
    print(f"输出目录: {args.output_dir}")
    if label_order:
        print(f"标签顺序: {label_order}")
    if label_merge_config:
        print(f"标签合并配置: {label_merge_config}")
    print("-" * 60)
    
    # 创建Mask生成器
    generator = MaskGenerator(label_order, label_merge_config)
    
    # 批量处理
    stats = generator.process_directory(args.input_dir, args.output_dir)
    
    # 输出统计信息
    print(f"\n{'='*60}")
    print("处理完成！")
    print(f"总文件数: {stats['total']}")
    print(f"成功处理: {stats['success']}")
    print(f"处理失败: {stats['failed']}")
    print(f"成功率: {stats['success']/stats['total']*100:.1f}%" if stats['total'] > 0 else "成功率: 0%")
    print("="*60)


if __name__ == "__main__":
    main()
